#include "coroutineHook.h"

static accept_t _accept = NULL;
static recv_t _recv = NULL;

ssize_t recv(int sockfd, void *buf, size_t len, int flags)
{
    Scheduler *s = GetInstance(); // 获取全局调度器
    Coroutine *co = s->coroutines[s->current_id];   
    // 设置非阻塞
    int flagsfc = fcntl(sockfd, F_GETFL, 0);
    fcntl(sockfd, F_SETFL, flagsfc | O_NONBLOCK);   
    // 调用原始recv
    while (1)
    {
        int clientfd = _recv(sockfd, (char *)buf, len, flags);
        if (clientfd < 0)
        {
            co->state = COROUTINE_SUSPEND;
            struct epoll_event ev;
            ev.data.fd = s->current_id;
            ev.events = EPOLLIN;
            epoll_ctl(s->epfd, EPOLL_CTL_ADD, co->fd, &ev);
            coroutine_yield();
            continue;
        }
        else if (clientfd == 0)
        {
            epoll_ctl(s->epfd, EPOLL_CTL_DEL, co->fd, NULL);
            close(co->fd);
            free(co);
            free(buf);
            break;
        }
        printf("recv:%s\n", (char *)buf);
    }
}

static void recv_cb(void *arg)
{
    Scheduler *s = GetInstance(); // 获取全局调度器
    Coroutine *co = s->coroutines[s->current_id];
    char *buf = (char *)malloc(1024);
    recv(co->fd, buf, 1024, 0);
}

// Hook后的accept函数
int accept(int sockfd, struct sockaddr *addr, socklen_t *addrlen)
{
    Scheduler *s = GetInstance(); // 获取全局调度器
    Coroutine *co = s->coroutines[s->current_id];

    // 设置非阻塞
    int flags = fcntl(sockfd, F_GETFL, 0);
    fcntl(sockfd, F_SETFL, flags | O_NONBLOCK);

    // 调用原始accept
    while (1)
    {
        int clientfd = _accept(sockfd, addr, addrlen);
        if (clientfd < 0)
        {
            co->state = COROUTINE_SUSPEND;
            struct epoll_event ev;
            ev.data.fd = s->current_id;
            ev.events = EPOLLIN;
            epoll_ctl(s->epfd, EPOLL_CTL_ADD, co->fd, &ev);
            // printf("add----co->fd:%d\n",co->fd);
            // printf("no connect,yield\n");
            coroutine_yield();
            continue;
        }
        int id = coroutine_create(clientfd, recv_cb, (void *)0);

        struct epoll_event ev;
        ev.data.fd = id;
        ev.events = EPOLLIN;
        epoll_ctl(s->epfd, EPOLL_CTL_ADD, clientfd, &ev);

        if (clientfd % 1000 == 0)
        {
            struct timeval current;
            gettimeofday(&current, NULL);
            int time_used = TIME_SUB_MS(current, begin);
            memcpy(&begin, &current, sizeof(struct timeval));
            printf("accept lientfd: %d\t time_used: %d\n", clientfd, time_used);
        }
    }
}

// 初始化Hook
void init_hooks()
{
    _accept = dlsym(RTLD_NEXT, "accept");
    _recv = dlsym(RTLD_NEXT, "recv");
}